package com.llmops.demo;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.openai.OpenAiChatModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.service.AiServices;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;

import java.util.List;

import static dev.langchain4j.data.document.loader.FileSystemDocumentLoader.loadDocuments;

    public class Easy_RAG_Example {


        interface Assistant {

            String chat(String userMessage);
        }

        public static void main(String[] args) {
            OpenAiChatModel chatModel = OpenAiChatModel.builder()
                    .baseUrl("https://api.deepseek.com/v1")
                    .apiKey("sk-0c3e6a9f9f7a4e63bb855290c544183c")
                    .modelName("deepseek-chat")
                    .build();

            // First, let's load documents that we want to use for RAG
            List<Document> documents = loadDocuments("D:\\新建文件夹\\llmops\\src\\main\\java\\com\\llmops\\demo\\ragdata");

            // Second, let's create an assistant that will have access to our documents
            Assistant assistant = AiServices.builder(Assistant.class)
                    .chatModel(chatModel) // it should use OpenAI LLM
                    .chatMemory(MessageWindowChatMemory.withMaxMessages(10)) // it should remember 10 latest messages
                    .contentRetriever(createContentRetriever(documents)) // it should have access to our documents
                    .build();

            String answer = assistant.chat("LangChain4j是什么");
            System.out.printf(answer);

        }

        private static ContentRetriever createContentRetriever(List<Document> documents) {

            // Here, we create an empty in-memory store for our documents and their embeddings.
            InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

            // Here, we are ingesting our documents into the store.
            // Under the hood, a lot of "magic" is happening, but we can ignore it for now.
            EmbeddingStoreIngestor.ingest(documents, embeddingStore);

            // Lastly, let's create a content retriever from an embedding store.
            return EmbeddingStoreContentRetriever.from(embeddingStore);
        }
    }
